# encoding: utf-8

from pathlib import Path
from pydantic import Field
from typing import Union
from loguru import logger
import shutil
import re


DOCKERFILE_TEMPLATE = """
FROM  {{python_version}}-slim-bookworm
COPY --from=ghcr.io/astral-sh/uv:latest /uv /uvx /bin/

# 设置工作目录
WORKDIR /app

# 复制依赖文件
COPY requirements.txt .

# 使用 uv 安装基础依赖到系统环境
RUN uv pip install --system -r requirements.txt -i https://pypi.tuna.tsinghua.edu.cn/simple

# 复制应用代码和启动脚本
COPY app/ ./app/
COPY start.sh .

# 创建依赖目录
RUN mkdir -p /dependencies

# 创建文件目录
RUN mkdir -p {{files_dir}}

# 设置启动脚本权限
RUN chmod +x start.sh

# 暴露端口
EXPOSE {{port}}

# 使用启动脚本替代直接的 uvicorn 命令
CMD ["./start.sh"]
"""

REGEX = re.compile(r"\{\{([a-zA-Z_][a-zA-Z0-9_]{0,29}|#histories#|#query#|#context#)\}\}")
WITH_VARIABLE_TMPL_REGEX = re.compile(
    r"\{\{([a-zA-Z_][a-zA-Z0-9_]{0,29}|#[a-zA-Z0-9_]{1,50}\.[a-zA-Z0-9_\.]{1,100}#|#histories#|#query#|#context#)\}\}"
)


class TemplateParser:
    """
    Rules:

    1. Template variables must be enclosed in `{{}}`.
    2. The template variable Key can only be: letters + numbers + underscore, with a maximum length of 16 characters,
       and can only start with letters and underscores.
    3. The template variable Key cannot contain new lines or spaces, and must comply with rule 2.
    4. In addition to the above, 3 types of special template variable Keys are accepted:
       `{{#histories#}}` `{{#query#}}` `{{#context#}}`. No other `{{##}}` template variables are allowed.
    """

    def __init__(self, template: str, with_variable_tmpl: bool = False):
        self.template = template
        self.with_variable_tmpl = with_variable_tmpl
        self.regex = WITH_VARIABLE_TMPL_REGEX if with_variable_tmpl else REGEX
        self.variable_keys = self.extract()

    def extract(self) -> list:
        # Regular expression to match the template rules
        return re.findall(self.regex, self.template)

    def format(self, inputs: dict, remove_template_variables: bool = True) -> str:
        def replacer(match):
            key = match.group(1)
            value = inputs.get(key, match.group(0))  # return original matched string if key not found

            if remove_template_variables and isinstance(value, str):
                return TemplateParser.remove_template_variables(value, self.with_variable_tmpl)
            return value

        prompt = re.sub(self.regex, replacer, self.template)
        return re.sub(r"<\|.*?\|>", "", prompt)

    @classmethod
    def remove_template_variables(cls, text: str, with_variable_tmpl: bool = False):
        return re.sub(WITH_VARIABLE_TMPL_REGEX if with_variable_tmpl else REGEX, r"{\1}", text)


class DockerSettings(object):

    def __init__(self):
        root_dir = Path(__file__).parent.parent
        self.sandbox_dir = root_dir / 'sandboxes' / 'python'
        self.python_version = 'python:3.12'
        self.files_dir = '/files'
        self.port = 8194
        self.dockerfile = TemplateParser(DOCKERFILE_TEMPLATE)

    def set_python_version(
            self, python_version: str = Field(default='python:3.12', description="python版本号.")
    ) -> None:
        self.python_version = python_version

    def set_requirements(
            self, requirement_file: Union[str, Path] = Field(default=..., description="python安装依赖文件.")
    ) -> None:
        if not Path(requirement_file).exists():
            logger.error(f"{requirement_file} not exists.")
            return
        target_file = (self.sandbox_dir / 'requirements.txt').as_posix()
        shutil.copy(requirement_file, target_file)

    def set_port(
            self, port: int = Field(default=8194, description='python代码沙箱服务端口号')
    ) -> None:
        self.port = port

    def set_file_dir(
            self, file_dir: Union[str, Path] = Field(default=..., description="python沙箱文件存放路径.")
    ) -> None:
        self.files_dir = file_dir

    def save(self):
        """保存Dockerfile"""
        file_content = self.dockerfile.format(
            inputs={
                "python_version": self.python_version,
                "files_dir": self.files_dir,
                "port": f'{self.port}'
            },
            remove_template_variables=True
        )
        docker_file = self.sandbox_dir / 'Dockerfile'
        with open(docker_file, 'w', encoding='utf-8', newline='') as fw:
            fw.write(file_content)


if __name__ == '__main__':
    settings = DockerSettings()
    settings.set_python_version('python:3.12')
    settings.set_file_dir('/files')
    settings.set_port(8194)
    settings.save()
